package com.gitee.wsl.func.filter.sample

import com.gitee.wsl.ext.base.Arrays.fill
import com.gitee.wsl.ext.base.SIZE
import com.gitee.wsl.func.filter.IFilterSet
import com.gitee.wsl.lang.bits.FixedBitSet
import kotlin.math.ceil
import kotlin.math.ln


/**
 * BloomKFilter is variation of [BloomFilter]. Unlike BloomFilter, BloomKFilter will spread 'k' hash bits within
 * same cache line for better L1 cache performance. The way it works is, First hash code is computed from key which is
 * used to locate the block offset (n-longs in bitset constitute a block) Subsequent 'k' hash codes are used to spread
 * hash bits within the block. By default block size is chosen as 8, which is to match cache line size (8 longs = 64
 * bytes = cache line size). Refer [BloomKFilter.add] for more info.
 *
 *
 * This implementation has much lesser L1 data cache misses than [BloomFilter].
 */
class BloomKFilter<T> : IFilterSet<T> {
    private val masks = LongArray(DEFAULT_BLOCK_SIZE)
    private val bitSet: FixedBitSet
    val bizSetSize: Int
    val k: Int
    private val hashFn: Hash64Fn<T>

    private var numberOfAddedElements = 0

    // spread k-1 bits to adjacent longs, default is 8
    // spreading hash bits within blockSize * longs will make bloom filter L1 cache friendly
    // default block size is set to 8 as most cache line sizes are 64 bytes and also AVX512 friendly
    private val totalBlockCount: Int

    constructor(numBits: Long, maxNumEntries: Long, hashFn: Hash64Fn<T>) {
        checkArgument(maxNumEntries > 0, "expectedEntries should be > 0")
        this.hashFn = hashFn
        // long numBits = optimalNumOfBits(maxNumEntries, DEFAULT_FPP);
        this.k = BloomFilter3.optimalNumOfHashFunctions(maxNumEntries, numBits)
        val nLongs = ceil(numBits.toDouble() / Long.SIZE.toDouble()).toInt()
        // additional bits to pad long array to block size
        val padLongs = DEFAULT_BLOCK_SIZE - nLongs % DEFAULT_BLOCK_SIZE
        this.bizSetSize = (nLongs + padLongs) * Long.SIZE
        this.bitSet = FixedBitSet(this.bizSetSize)
        checkArgument((bitSet.dataLength % DEFAULT_BLOCK_SIZE) == 0, "bitSet has to be block aligned")
        this.totalBlockCount = bitSet.dataLength / DEFAULT_BLOCK_SIZE
    }

    /**
     * A constructor to support rebuilding the BloomFilter from a serialized representation.
     *
     * @param bits
     * @param numFuncs
     */
    constructor(bits: LongArray, numFuncs: Int, numberOfAddedElements: Int, hashFn: Hash64Fn<T>) {
        this.hashFn = hashFn
        this.numberOfAddedElements = numberOfAddedElements
        bitSet = FixedBitSet(bits)
        this.bizSetSize = bits.size * Long.SIZE
        this.k = numFuncs
        checkArgument((bitSet.dataLength % DEFAULT_BLOCK_SIZE) == 0, "bitSet has to be block aligned")
        this.totalBlockCount = bitSet.dataLength / DEFAULT_BLOCK_SIZE
    }

    override fun add(elm: T) {
        // We use the trick mentioned in "Less Hashing, Same Performance: Building a Better Bloom Filter"
        // by Kirsch et.al. From abstract 'only two hash functions are necessary to effectively
        // implement a Bloom filter without any loss in the asymptotic false positive probability'

        // Lets split up 64-bit hashcode into two 32-bit hash codes and employ the technique mentioned
        // in the above paper

        val hash64: Long = hashFn(elm)
        addHash(hash64)
    }

    private fun addHash(hash64: Long) {
        val hash1 = hash64.toInt()
        val hash2 = (hash64 ushr 32).toInt()

        var firstHash = hash1 + hash2
        // hashcode should be positive, flip all the bits if it's negative
        if (firstHash < 0) {
            firstHash = firstHash.inv()
        }

        // first hash is used to locate start of the block (blockBaseOffset)
        // subsequent K hashes are used to generate K bits within a block of words
        val blockIdx = firstHash % totalBlockCount
        val blockBaseOffset = blockIdx shl DEFAULT_BLOCK_SIZE_BITS
        val data = bitSet.data
        for (i in 1..k) {
            var combinedHash = hash1 + ((i + 1) * hash2)
            // hashcode should be positive, flip all the bits if it's negative
            if (combinedHash < 0) {
                combinedHash = combinedHash.inv()
            }
            // LSB 3 bits is used to locate offset within the block
            val absOffset = blockBaseOffset + (combinedHash and DEFAULT_BLOCK_OFFSET_MASK)
            // Next 6 bits are used to locate offset within a long/word
            val bitPos = (combinedHash ushr DEFAULT_BLOCK_SIZE_BITS) and DEFAULT_BIT_OFFSET_MASK
            data[absOffset] = data[absOffset] or (1L shl bitPos)
        }
    }

    override fun filter(elm: T): Boolean {
        val hash64: Long = hashFn(elm)
        return containsHash(hash64)
    }

    private fun containsHash(hash64: Long): Boolean {
        val hash1 = hash64.toInt()
        val hash2 = (hash64 ushr 32).toInt()

        var firstHash = hash1 + hash2
        // hashcode should be positive, flip all the bits if it's negative
        if (firstHash < 0) {
            firstHash = firstHash.inv()
        }

        // first hash is used to locate start of the block (blockBaseOffset)
        // subsequent K hashes are used to generate K bits within a block of words
        // To avoid branches during probe, a separate masks array is used for each longs/words within a block.
        // data array and masks array are then traversed together and checked for corresponding set bits.
        val blockIdx = firstHash % totalBlockCount
        val blockBaseOffset = blockIdx shl DEFAULT_BLOCK_SIZE_BITS

        // iterate and update masks array
        for (i in 1..k) {
            var combinedHash = hash1 + ((i + 1) * hash2)
            // hashcode should be positive, flip all the bits if it's negative
            if (combinedHash < 0) {
                combinedHash = combinedHash.inv()
            }
            // LSB 3 bits is used to locate offset within the block
            val wordOffset = combinedHash and DEFAULT_BLOCK_OFFSET_MASK
            // Next 6 bits are used to locate offset within a long/word
            val bitPos = (combinedHash ushr DEFAULT_BLOCK_SIZE_BITS) and DEFAULT_BIT_OFFSET_MASK
            masks[wordOffset] = masks[wordOffset] or (1L shl bitPos)
        }

        val data = bitSet.data
        // traverse data and masks array together, check for set bits
        var expected: Long = 0
        for (i in 0..<DEFAULT_BLOCK_SIZE) {
            val mask = masks[i]
            expected = expected or ((data[blockBaseOffset + i] and mask) xor mask)
        }

        // clear the mask for array reuse (this is to avoid masks array allocation in inner loop)
        fill(masks, 0)

        // if all bits are set, expected should be 0
        return expected == 0L
    }

    /**
     * Sets all bits to false in the Bloom filter.
     */
    override fun clear() {
        bitSet.clear()
        numberOfAddedElements = 0
    }

    override val isEmpty: Boolean
        get() = numberOfAddedElements == 0

    fun count(): Int {
        return numberOfAddedElements
    }

    companion object {
        // static final float DEFAULT_FPP = 0.05f;
        private const val DEFAULT_BLOCK_SIZE = 8
        private val DEFAULT_BLOCK_SIZE_BITS = (ln(DEFAULT_BLOCK_SIZE.toDouble()) / ln(2.0)).toInt()
        private val DEFAULT_BLOCK_OFFSET_MASK = DEFAULT_BLOCK_SIZE - 1
        private val DEFAULT_BIT_OFFSET_MASK: Int = Long.SIZE - 1
        fun checkArgument(expression: Boolean, message: String) {
            require(expression) { message }
        }

//        fun forString(inputEntries: Long, fpp: Double): BloomKFilter<String?> {
//            val bitSetSize: Int = BloomFilter3.optimalNumOfBits(inputEntries, fpp)
//            return BloomKFilter<String?>(
//                bitSetSize.toLong(),
//                inputEntries.toInt().toLong(),
//                HashHelper.murmur3_64_string()
//            )
//        }
    }
}
